#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>

/* TODO: fix include path */
#include "../include/main.h"
#include "../include/menu.h"
#include "../include/crypto.h"
#include "../source/imageADT/imageADT.h"

static int cryptDes(appInit *params);
static int cryptAes(appInit *params);

static int getAlgorithm(appInit *params);
static int getOp(appInit *params, int algorithm);
static int getMethod(appInit *params);
static int checkKeyIvLen(appInit *params);


int
main(int argc, char **argv)
{
    appInit params;
    char *salida;
    int algorithm;
	int ret;

    getParameters(argc, argv, &params);

    if (!validateParameters(&params) )
    {
        usage(argv[0]);
        return 0;
    }

    if ( !checkKeyIvLen(&params) ) {
        fprintf(stderr, "iv length or k length is wrong.\n");
        fprintf(stderr, "They must be of length 8 for DES and len 16 for AES.\n");
        return 0;
    }

    algorithm = getAlgorithm(&params);

    switch(params.op)
    {
        case 'e':
            printf("Encrypting...\n");

	    if ( algorithm == DES )
			ret = cryptDes(&params);
	    else
			ret =  cryptAes(&params);

			if ( ret == 1 )
				printf("Encryption finished.\n");
			else
				printf("Encryption failed\n");
            break;

        case 'd':
            printf("Decrypting...\n");

            if( algorithm == DES )
				ret = cryptDes(&params);
			else
				ret = cryptAes(&params);

			if ( ret == 1 )
				printf("Decryption finished.\n");
			else
				printf("Decryption failed\n");
            break;

        default:
            fprintf(stderr, "The selected operation doesn't exist.\n");
            break;
    }

    freeParameters(&params);
    return 0;
}

static int
cryptDes(appInit *params) 
{
	int op;
	int method;
	imageADT image;
	FILE *in;
	FILE *out; 
	DES_key_schedule *ks;
	DES_cblock *iv;
	char buffer[DES_BLOCK_SIZE];
	char *bufferCrypt;

	if( params == NULL )
		return 0;

	op = getOp( params, DES );
	method = getMethod( params );

	if( (image = openImage(params->inName)) == NULL ) {
        fprintf(stderr, "Couldn't open the image.\n");
        return 0;
    }

	if( (in = getFile(image)) == NULL || !isMod(image, 8) || !headerCheck(image) ) {
        fprintf(stderr, "Couldn't open the image.\n");
        fprintf(stderr, "Check if the path is correct.\n");
        fprintf(stderr, "Check if it's a bmp file.\n");
        fprintf(stderr, "Check if the size of the images is mod 8.\n");
		closeImage(image);
		return 0;
	}

	if( (out = createOutputFile(image, params->outName)) == NULL ){
        fprintf(stderr, "Couldn't create the output file.\nNo space left?\n");
		closeImage(image);
		return 0;
	}

	if( (ks = malloc(sizeof(DES_key_schedule))) == NULL ){
		closeImage(image);
		fclose(out);
		return 0;
	}
	
	if( (iv = malloc(sizeof(DES_cblock))) == NULL ){
		closeImage(image);
		fclose(out);
		free(ks);
		return 0;
	}

 	if( !initDes( method, params->password, params->key, params->iv, ks, iv ) ){
		closeImage(image);
		fclose(out);
		free(ks);
		free(iv);
        fprintf(stderr, "Couldn't init des, check your password/key-iv\n");
		return 0;
	}	

	while ( fread( buffer, 1, DES_BLOCK_SIZE, in ) == DES_BLOCK_SIZE ) {

			bufferCrypt = des( op, method, ks, iv, buffer );
	
			if ( fwrite(bufferCrypt, 1, DES_BLOCK_SIZE, out) != DES_BLOCK_SIZE ){
				fprintf(stderr, "ERROR: Couldn't write the outfile.\n");
				closeImage(image);
				fclose(out);
				free(ks);
				free(iv);
				free(bufferCrypt);
				return 0;
			}

			free(bufferCrypt);
	}

	closeImage(image);
	fclose(out);
	free(ks);
	free(iv);

	return 1;
}

static int
cryptAes(appInit *params) 
{
	int op;
	int method;
	imageADT image;
	FILE *in;
	FILE *out; 
	AES_KEY *aesKey;
	unsigned char *iv;
	char buffer[AES_BLOCK_SIZE];
	char *bufferCrypt;

	if( params == NULL )
		return 0;

	op = getOp( params, AES );
	method = getMethod( params );

	if( (image = openImage(params->inName)) == NULL ) {
        fprintf(stderr, "Couldn't open the image.\n");
        return 0;
    }

    if( (in = getFile(image)) == NULL || !isMod(image, 16) || !headerCheck(image) ) {
        fprintf(stderr, "Couldn't open the image.\n");
        fprintf(stderr, "Check if the path is correct.\n");
        fprintf(stderr, "Check if it's a bmp file.\n");
        fprintf(stderr, "Check if the size of the images is mod 16.\n");
		closeImage(image);
		return 0;
	}

	if( (out = createOutputFile(image, params->outName)) == NULL ){
        fprintf(stderr, "Couldn't create the output file.\nNo space left?\n");
		closeImage(image);
		return 0;
	}

	if( (aesKey = malloc(sizeof(AES_KEY))) == NULL ){
		closeImage(image);
		fclose(out);
		return 0;
	}
	
	if( (iv = malloc(AES_KEY_SIZE)) == NULL ){
		closeImage(image);
		fclose(out);
		free(aesKey);
		return 0;
	}

 	if( !initAes( op, method, params->password, params->key, params->iv, aesKey, iv ) ){
        fprintf(stderr, "Couldn't init des, check your password/key-iv\n");
		closeImage(image);
		fclose(out);
		free(aesKey);
		free(iv);
		return 0;
	}	

	while ( fread( buffer, 1, AES_BLOCK_SIZE, in ) == AES_BLOCK_SIZE ) {

			bufferCrypt = aes( op, method, aesKey, iv, buffer );
	
			if ( fwrite(bufferCrypt, 1, AES_BLOCK_SIZE, out) != AES_BLOCK_SIZE ){
				printf("ERROR: Couldn't write the outfile.\n");
				closeImage(image);
				fclose(out);
				free(aesKey);
				free(iv);
				free(bufferCrypt);
				return 0;
			}

			free(bufferCrypt);
	}

	closeImage(image);
	fclose(out);
	free(aesKey);
	free(iv);

	return 1;
}

static int
getAlgorithm(appInit *params)
{
	if( strcmp(params->algorithm,"des") == 0 )
		return DES;
	else
		return AES;
}

static int
getOp(appInit *params, int algorithm)
{
	if( params->op == 'e' ){
		if( algorithm == DES )
			return DES_ENCRYPT;
		else
			return AES_ENCRYPT;
	}
	else{
		if( algorithm == DES )
			return DES_DECRYPT;
		else
			return AES_DECRYPT;
	}

}

static int
checkKeyIvLen(appInit *params) {
    if ( params->password != NULL )
        return 1;

    if ( strcmp(params->algorithm, "des") == 0 ) {
        if ( strlen(params->key) != 8 || strlen(params->iv) != 8 )
            return 0;
    } else {
        if ( strlen(params->key) != 16 || strlen(params->iv) != 16 )
            return 0;
    }

    return 1;
}

static int
getMethod(appInit *params)
{
	if( strcmp(params->method,"ecb") == 0 )
		return ECB;
	else if( strcmp(params->method,"cfb") == 0 )
		return CFB;
	else if( strcmp(params->method,"ofb") == 0 )
		return OFB;
	else if( strcmp(params->method,"cbc") == 0 )
		return CBC;
}



